from fastapi import APIRouter
from logger.logger import ServiceLogger
from server.preprocessing.image import ImagePreprocessor

from server.inference.client import TritonInferenceClient
from typing import Optional
from fastapi import HTTPException
from starlette.responses import JSONResponse
import requests
import yaml
from pathlib import Path

# fastapi接受请求, 调用preprocessing, inference, postprocessing, 返回结果

flashlightRouter=APIRouter()
imagePreprocessor = ImagePreprocessor()
logger = ServiceLogger()

def read_image_from_url(url: str) -> bytes:
    response = requests.get(url)
    if response.status_code != 200:
        raise HTTPException(status_code=400, detail="Failed to fetch image from URL")
    return response.content

def load_config(config_path: str) -> dict:
    with open(Path(config_path), "r") as file:
        return yaml.safe_load(file)

@flashlightRouter.get("/flashlight/", tags=["flashlight"])
async def identify_images(fileurl: str, model: Optional[str] = None):
    # 判断验证参数是否合法
    if not fileurl:
        logger.log_message("Missing required parameter: fileurl")
        # JSON 返回参数错误
        return JSONResponse(status=400, content={"message": "Missing required parameter: fileurl"})

    # 对图片资源的前置处理
    ServiceLogger.log_message("Starting flashlight recognition")
    image_bytes = read_image_from_url(fileurl)
    ndarray = imagePreprocessor.preprocess(image_bytes)

    # 读取 config.interence.yaml 中的配置
    config = load_config("config.inference.yaml")

    # 推送请求到 triton
    client = TritonInferenceClient(
        server_url=config.inference.triton_server.url,
        model_name=model if model else config.inference.model_name,
        input_name="input0",
        output_name="output0"
    )
    inputs = {"input0": ndarray}
    outputs = await client.run_inference(inputs)
    print(outputs)

    await client.close()

    # 返回处理结果（这里假设返回一个简单的字典表示识别结果）
    return {"success": True, "outputs": outputs}

